#!/usr/bin/env python3
"""
Verify that code_content has been successfully added to the graph
"""

import pickle
from collections import defaultdict

# Load the modified graph
pkl_file = "/data_ext/ref_code/LocAgent/index_data_strict/pp_data/graph_index_v2.3/business-payout-with-code-content.pkl"

print("Loading modified graph...")
with open(pkl_file, "rb") as f:
    G = pickle.load(f)

print(f"\n=== Graph Statistics ===")
print(f"Total nodes: {G.number_of_nodes()}")
print(f"Total edges: {G.number_of_edges()}")

# Check code_content attribute
stats = {
    'has_code': 0,
    'has_code_content': 0,
    'has_both': 0,
    'has_neither': 0
}

node_types_with_code_content = defaultdict(int)

for node_id in G.nodes():
    node_data = G.nodes[node_id]
    has_code = 'code' in node_data
    has_code_content = 'code_content' in node_data

    if has_code and has_code_content:
        stats['has_both'] += 1
        if 'type' in node_data:
            node_types_with_code_content[node_data['type']] += 1
    elif has_code:
        stats['has_code'] += 1
    elif has_code_content:
        stats['has_code_content'] += 1
    else:
        stats['has_neither'] += 1

print(f"\n=== Code Attribute Statistics ===")
print(f"Nodes with both 'code' AND 'code_content': {stats['has_both']}")
print(f"Nodes with only 'code': {stats['has_code']}")
print(f"Nodes with only 'code_content': {stats['has_code_content']}")
print(f"Nodes with neither: {stats['has_neither']}")

print(f"\n=== Node Types with code_content ===")
for node_type, count in sorted(node_types_with_code_content.items()):
    print(f"{node_type}: {count} nodes")

# Show some examples of nodes with code_content
print(f"\n=== Sample Nodes with code_content ===")
samples_shown = 0
for node_id in G.nodes():
    node_data = G.nodes[node_id]
    if 'code_content' in node_data and node_data['code_content']:
        print(f"\nNode: {node_id}")
        print(f"  Type: {node_data.get('type', 'unknown')}")

        code_lines = node_data['code_content'].split('\n')
        print(f"  Code content: {len(code_lines)} lines")

        # Show first 3 lines
        print(f"  First 3 lines:")
        for i, line in enumerate(code_lines[:3], 1):
            if len(line) > 80:
                line = line[:80] + "..."
            print(f"    {i}: {line}")

        samples_shown += 1
        if samples_shown >= 3:
            break

print(f"\n✅ Verification complete!")
if stats['has_both'] > 0:
    print(f"Successfully added code_content to {stats['has_both']} nodes!")
else:
    print("⚠️ No nodes have code_content attribute!")